import os
from IMAGE_QUERY.extract_cnn_vgg16_keras import VGGNet
from pymongo import MongoClient


def get_imlist(path):
    return [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.jpg')]


# try:
#     conn = MongoClient()
#     print("Connected successfully!!!")
# except:
#     print("Could not connect to MongoDB")

# database and table
# db = conn.toutiao
# collection = db.toutiao_feats_04

'''
 Extract features and index the images
'''
if __name__ == "__main__":

    # the root path of images
    db = 'C:/Users/Administrator/Desktop'
    img_list = get_imlist(db)

    print("--------------------------------------------------")
    print("feature extraction starts......")
    print("--------------------------------------------------")

    items = []
    model = VGGNet()
    for i, img_path in enumerate(img_list):
        feat = model.extract_feat(img_path)
        img_name = os.path.split(img_path)[1]
        print("extracting feature from image No. %d , %d images in total" % ((i + 1), len(img_list)))

        item = {'img_name': img_name, 'feats': feat.tolist()}
        # items.append(item)
        print(item)
    # collection.insert_many(items)
    # conn.close()

    print("--------------------------------------------------")
    print("        writing feature extraction results .......")
    print("--------------------------------------------------")

    print("write over")

